#!/usr/bin/env python
# coding: utf-8

# In[1]:


# -*- coding: utf-8 -*-
"""
数据划分模块
输入：encoded_data.csv
输出：train_set.csv 和 test_set.csv
"""

import pandas as pd
import os  
from sklearn.model_selection import train_test_split

# 配置参数
SEED = 2025
TEST_SIZE = 0.3
DATA_PATH = 'encoded_data.csv'
TARGET_COL = 'label'  

def main():
    # 1. 加载预处理数据
    try:
        preprocessed_data = pd.read_csv(DATA_PATH)
        print(f"数据加载成功，维度：{preprocessed_data.shape}")
    except FileNotFoundError:
        print(f"错误：找不到文件 {DATA_PATH}")
        return

    # 2. 验证目标列存在性
    if TARGET_COL not in preprocessed_data.columns:
        print(f"错误：目标列 '{TARGET_COL}' 不存在，可用列：{preprocessed_data.columns.tolist()}")
        return

    # 3. 分离特征和标签
    X = preprocessed_data.drop(columns=[TARGET_COL])
    y = preprocessed_data[TARGET_COL]

    # 4. 分层划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y,
        test_size=TEST_SIZE,
        random_state=SEED,
        stratify=y
    )

    # 5. 合并并保存数据集
    pd.concat([X_train, y_train], axis=1).to_csv('train_set.csv', index=False)
    pd.concat([X_test, y_test], axis=1).to_csv('test_set.csv', index=False)

    # 6. 打印结果
    print("\n划分结果：")
    print(f"训练集路径：{os.path.abspath('train_set.csv')}")  
    print(f"测试集路径：{os.path.abspath('test_set.csv')}")

if __name__ == '__main__':
    main()


# In[2]:


#测试划分后的训练集和测试集的正负标签个数

import pandas as pd
import os

def check_label_distribution(file_path, target_col='label'):
    """检查数据集的标签分布"""
    try:
        df = pd.read_csv(file_path)
        print(f"数据集路径：{os.path.abspath(file_path)}")
        print(f"总样本数：{len(df)}")

        # 统计标签分布
        label_counts = df[target_col].value_counts().sort_index()
        ratios = df[target_col].value_counts(normalize=True).sort_index()

        # 格式化输出
        print("标签分布统计：")
        for label, count in label_counts.items():
            print(f"Label {label}: {count:5d} 个样本 | 占比 {ratios[label]:.2%}")
        print("-" * 50)

    except FileNotFoundError:
        print(f"错误：文件 {file_path} 不存在")
    except KeyError:
        print(f"错误：目标列 {target_col} 不存在，可用列：{df.columns.tolist()}")

# 检查三个数据集的分布
print("="*50 + "\n原始数据集分布")
check_label_distribution('processed_data_cleaned.csv')

print("\n" + "="*50 + "\n训练集分布")
check_label_distribution('train_set.csv')

print("\n" + "="*50 + "\n测试集分布")
check_label_distribution('test_set.csv')


# In[ ]:




